package SSF.App.DDoS;

import SSF.OS.NetFlow.*;
import SSF.Util.Streams.*;
import SSF.Net.Util.*;
import java.io.*;
import java.util.*;

/** Cross Domain DDoS Tracer */
public class CD_DDoSTracer extends BasicPlayer
{
    /** some constants*/
    public static final int SET_QUERY_FAIL = -1;
    public static final int SET_QUERY_SUCC = 0;

    /** store information of related ingress traffic*/
    protected Hashtable ingressTable;

    /** store inforamtion of related egress traffic */
    protected Hashtable egressTable;

    /** the target IP address */
    protected int targetIP;

    /** start time */
    protected float startTime=0;

    /** end time */
    protected float endTime=Float.MAX_VALUE;

    /** statistic data of a domain */
    long totalIngressSize = 0;
    long totalIngressPkts = 0;
    long totalIngressFlows = 0;
    long totalEgressSize = 0;
    long totalEgressPkts = 0;
    long totalEgressFlows = 0;

    private int netFlowTypeCode=-1;

    public CD_DDoSTracer(String streamID)
    {
	super(streamID);
	ingressTable = new Hashtable();
	egressTable = new Hashtable();
    }

    /**
     * reset all the statistic values, clear the tables.
     */
    public void reset()
    {
	ingressTable.clear();
	egressTable.clear();
	totalIngressSize = 0;
	totalIngressPkts = 0;
	totalIngressFlows = 0;
	totalEgressSize = 0;
	totalEgressPkts = 0;
	totalEgressFlows = 0;
    }
    /**
     * This is called to set the parameters for query.
     */
    public int setQuery(int targetIP, float startTime, float endTime)
    {
	this.targetIP = targetIP;
	this.startTime = startTime;
	this.endTime = endTime;

	if (startTime >= endTime) {
	    System.err.println("Invalid startTime/endTime.");
	    return SET_QUERY_FAIL;
	}
	return SET_QUERY_SUCC;
    }

    public int receive(int type_code, int source_code, double timestamp, 
		       byte[] bytes, int offset, int length)
    {
	//check type
	if (netFlowTypeCode == -1) {
	    netFlowTypeCode = getRecordTypeCode("SSF.OS.NetFlow");
	}
	if (type_code != netFlowTypeCode) {
	    return 0;             // not a netflow record
	}

	// construct flow
	IpNetFlowWD flow = new IpNetFlowWD();
	flow.readFromBytes(bytes, length);

	// check the destination IP
	if (flow.dstAddr != targetIP) {
	    return 0;             // not related record.
	}

	// check whether it's within the time period
	if (flow.Last < startTime || flow.First > endTime) {
	    return 0;             // out of the time window.
	}

	// update ingress table if it's ingress related record/
	if (flow.inputType == IpNetFlowWD.EXTERNAL) {
	    totalIngressFlows++;
	    totalIngressPkts += flow.dPkts;
	    totalIngressSize += flow.dOctets;

	    String ikey = flow.nhi+" "+flow.input;
	    TraceInfo tinfo = (TraceInfo)ingressTable.get(ikey);
	    if (tinfo == null) {
		// new ingress link.
		tinfo = new TraceInfo(flow.nhi, flow.input, (long)flow.dOctets,
				      (long)flow.dPkts);
		ingressTable.put(ikey, tinfo);
	    } else {
		// update ingress info
		tinfo.update((long)flow.dOctets, (long)flow.dPkts);
	    }
	}

	// update egress table if it's egress related record..
	if (flow.outputType == IpNetFlowWD.EXTERNAL) {
	    totalEgressFlows++;
	    totalEgressPkts += flow.dPkts;
	    totalEgressSize += flow.dOctets;

	    String ekey = flow.nhi+" "+flow.output;
	    TraceInfo tinfo = (TraceInfo)egressTable.get(ekey);
	    if (tinfo == null) {
		// new ingress link.
		tinfo = new TraceInfo(flow.nhi, flow.output, 
				      (long)flow.dOctets,(long)flow.dPkts);
		egressTable.put(ekey, tinfo);
	    } else {
		// update ingress info
		tinfo.update((long)flow.dOctets, (long)flow.dPkts);
	    }
	}
	return 0;
    }

    /**
     * print out the content and a summary of ingress/egress table
     */
    public void printSummary()
    {
	System.out.println("--------- Ingress information ---------");
	TraceInfo tinfo;
	for (Enumeration enumVar = ingressTable.elements();
	     enumVar.hasMoreElements();)
	    System.out.println(((TraceInfo)enumVar.nextElement()).toString());

	System.out.println("--------- Egress information ----------");
	for (Enumeration enumVar = egressTable.elements();
	     enumVar.hasMoreElements();)
	    System.out.println((TraceInfo)enumVar.nextElement());

	System.out.println("\n----------==== Summarry ====----------");
	System.out.println("Destination IP: "+IP_s.IPtoString(targetIP)+"\n"
			   +"Start time: "+startTime
			   +"    End time: "+endTime+"\n"
			   +"total Ingress Pkts: "+totalIngressPkts+"\n"
			   +"total Egress Pkts: "+totalEgressPkts+"\n"
			   +"\tPkts difference: "
			   +(totalEgressPkts-totalIngressPkts)+"\n\n"
			   +"total Ingress Size: "+totalIngressSize+"\n"
			   +"total Egress Size: "+totalEgressSize+"\n"
			   +"\tSize difference: "
			   +(totalEgressSize-totalIngressSize)+"\n\n"
			   +"total Ingress Flows: "+totalIngressFlows+"\n"
			   +"total Egress Flows: "+totalEgressFlows+"\n"
			   +"\tFlows difference: "
			   +(totalEgressFlows-totalIngressFlows)
			   +"\n---------------------------------------");
    }


    public static void main(String[] args)
    {
	float startTime = 0;
	float endTime = Float.MAX_VALUE;
	String srcFileName=null;
	int targetIP=0;
	String streamID=null;

	// get argument.
	try {
	    switch (args.length) {
	    case 6:
		if (!args[0].equals("-t"))
		    usageError();
		startTime = Float.valueOf(args[1]).floatValue();
		endTime = Float.valueOf(args[2]).floatValue();
		targetIP = IP_s.StrToInt(args[3]);
		srcFileName = args[4];
		streamID = args[5];
		break;
	    case 3:
		targetIP = IP_s.StrToInt(args[0]);
		srcFileName = args[1];
		streamID = args[2];
		break;
	    default:
		usageError();
	    }
	} catch (Exception any) {
	    System.err.println("Error in parameter.");
	    System.err.println(any.getMessage());
	    usageError();
	}

	//for debug
	System.out.println("parameter check passed.");
	CD_DDoSTracer dosTracer = new CD_DDoSTracer(streamID);
	dosTracer.setQuery(targetIP, startTime, endTime);
	try {
	    dosTracer.connectRead(srcFileName);
	} catch (streamException se) {
	    System.err.println(se.getMessage());
	    se.printStackTrace();
	    System.exit(-1);
	}
	dosTracer.printSummary();
    }

    private static void usageError()
    {
	System.err.println("usage: java CD_DDoSTracer [-t startTime endTime] <target IP> <netflow file> <streamID>");
	System.exit(-1);
    }

    public Enumeration ingressInfo() 
    {
	return ingressTable.elements();
    }

    public Enumeration egressInfo()
    {
	return egressTable.elements();
    }
}
